import json

from typing import Generator, List

from fastapi.encoders import jsonable_encoder
from pydantic.tools import parse_obj_as
from sqlalchemy.orm import Session

from app.api import deps
from app.common.dateutils import DateUtils
from app.common.page import PageParam
from app.core.celery_app import celery_app
from app.core.constant_code import ConstantItem
from app.core.exception_code import ExceptionCode
from app.dao.bar_dao import barDao
from app.dao.instrument_dao import instrumentDao
from app.dao.picker_dao import pickerDao
from app.dao.strategy_code_dao import strategyCodeDao
from app.dao.strategy_dao import strategyDao
from app.db.redis_session import redisClient
from app.manager.redis_manager import redisManager
from app.models.bar import Bar
from app.models.instrument import Instrument
from app.models.picker_instrument import PickerInstrument
from app.schemas.bar import CandlestickQuery
from app.schemas.picker import CeleryTaskMsg, PickerTask, PublicConditionParam, StrategyConditionParam, PickerBase
from app.schemas.strategy_code import StrategyCodeQuery
from assertpy.assertions import assert_that
# from ifcompiler.executor.tdxexecutor import TdxExecutor
from ifcompiler.drawing.drawing import Candlestick
from ifcompiler.compiler.compiler import TdxCompiler
from ifcompiler.context import CompilerContext
from iflib.platforms.tdx import CLOSE, HIGH, OPEN, LOW, VOL, AMOUNT, DATETIME
import numpy as np


@celery_app.task(acks_late=True)
def picker(msg: str) -> str:
    db: Session = next(deps.getDb())
    taskMsg = CeleryTaskMsg(**json.loads(msg))
    taskMsgBody = PickerTask(**taskMsg.msgBody)
    picker = pickerDao.selectById(db, taskMsgBody.pickerId)
    # assert_that(picker.userId).is_not_equal_to(id).then_throw(ExceptionCode.)
    publicConditions = PublicConditionParam(**json.loads(picker.publicConditions))
    strategyConditions = parse_obj_as(List[StrategyConditionParam], json.loads(picker.strategyConditions))
    # 合并策略，id,参数，选择日期一致
    mergedStrategyConditions = []
    present = None
    while (len(strategyConditions) > 0):
        present = strategyConditions[0]
        mergedStrategyConditions.append(present)
        strategyConditions.remove(present)
        temp = []
        for s in strategyConditions:
            if present.strategyId == s.strategyId \
                    and present.pickDate == s.pickDate \
                    and json.dumps(jsonable_encoder(present.strategyParams)) == json.dumps(
                jsonable_encoder(s.strategyParams)):
                temp.append(s)
                present.expressions.extend(s.expressions)
        # 找到一致的,添加完表达式就删除
        for t in temp:
            strategyConditions.remove(t)
    strategyIds = set([s.strategyId for s in mergedStrategyConditions])
    compilers = {}
    for strategyId in strategyIds:
        # 查询策略
        strategy = strategyDao.selectById(db, strategyId)
        assert_that(strategy).is_none().then_throw(ExceptionCode.STRATEGY_DOES_NOT_EXIST)
        strategyCodeQuery = StrategyCodeQuery(state=ConstantItem.STATE_ENABLED, strategyId=strategyId)
        strategyCode = strategyCodeDao.selectOne(db, strategyCodeQuery)
        assert_that(strategyCode).is_none().then_throw(ExceptionCode.STRATEGY_DOES_NOT_EXIST)
        context = CompilerContext(strategyCode.code)
        # context.strategyVarsCode = variables
        # context.pickResultCode = expressions
        compilers[strategyId] = TdxCompiler(context)
    # 设置K线查询参数
    adjust = ConstantItem.ADJUST_FORE if publicConditions.isAdjustFore else ConstantItem.ADJUST_NO
    isSt = 0 if publicConditions.isRemoveSt else None
    tradeStatus = 1 if publicConditions.isRemoveUntrade else None
    # pickDate最大日期不是st的股票
    pageParam = PageParam(pageNum=1, pageSize=500)
    for instrumentId in taskMsgBody.instrumentIds:
        try:
            tempResult = []
            for cond in mergedStrategyConditions:
                # 装载数据
                candlestickQuery = CandlestickQuery(instrumentId=instrumentId,
                                                    frequency=picker.frequency,
                                                    adjust=adjust,
                                                    tradeStatus=tradeStatus,
                                                    barDate=cond.pickDate)
                dbData = barDao.getCandlestick(db, candlestickQuery, pageParam)
                dbData.reverse()
                # 装载数据
                candlestick = Candlestick()
                for row in dbData:
                    candlestick.closeSeries.append(row.close)
                    candlestick.openSeries.append(row.open)
                    candlestick.highSeries.append(row.high)
                    candlestick.lowSeries.append(row.low)
                    candlestick.dateSeries.append(DateUtils.dateToTimestamp13(row.barDate))
                    candlestick.volumeSeries.append(row.volume)
                    candlestick.amountSeries.append(row.amount)
                    candlestick.preCloseSeries.append(row.preClose)
                CLOSE.series = np.array(candlestick.closeSeries).astype(np.float)
                HIGH.series = np.array(candlestick.highSeries).astype(np.float)
                OPEN.series = np.array(candlestick.openSeries).astype(np.float)
                LOW.series = np.array(candlestick.lowSeries).astype(np.float)
                VOL.series = np.array(candlestick.volumeSeries).astype(np.float)
                AMOUNT.series = np.array(candlestick.amountSeries).astype(np.float)
                DATETIME.series = np.array(candlestick.dateSeries)
                # 初始化执行器
                # 转换策略参数
                variables = ""
                for param in cond.strategyParams:
                    variables += "%s = %s;\n" % (param.variable, param.value)
                # 转换选股结果表达式
                expressions = ""
                for expr in cond.expressions:
                    expressions += f" & {expr.parsedExpr}" if expressions else expr.parsedExpr

                # 编译源码
                compiler = compilers.get(cond.strategyId)
                context = compiler.context
                context.strategyVarsCode = variables
                context.pickResultCode = f'{context.varResult} = {expressions};'
                compiler.analyze().parse().generate().execute()
                # 加载数据
                tempResult.append(context.executeResult[context.varResult].series[-1])
                # drawings = executor.get_drawings()
            result = None
            for r in tempResult:
                if not result:
                    result = r
                else:
                    result &= r
            # 保存
            pickerDao.updatePickResult(db, picker.id, result)
            if result:
                pickerInstrument = PickerInstrument(pickerId=picker.id,
                                                    instrumentId=instrumentId,
                                                    probability=1)
                db.add(pickerInstrument)
            db.commit()
        except Exception as ex:
            print(f"{instrumentId}执行失败")
            pickerDao.updatePickResult(db, picker.id, False)
            db.commit()
            print(ex)
    db.close()
    return f"success"


@celery_app.task(acks_late=True)
def realtime_picker(msg: str) -> str:
    db: Session = next(deps.getDb())
    taskMsg = CeleryTaskMsg(**json.loads(msg))
    taskMsgBody = PickerTask(**taskMsg.msgBody)
    picker = pickerDao.selectById(db, taskMsgBody.pickerId)
    # assert_that(picker.userId).is_not_equal_to(id).then_throw(ExceptionCode.)
    publicConditions = PublicConditionParam(**json.loads(picker.publicConditions))
    strategyConditions = parse_obj_as(List[StrategyConditionParam], json.loads(picker.strategyConditions))
    # 合并策略，id,参数，选择日期一致
    mergedStrategyConditions = []
    present = None
    while (len(strategyConditions) > 0):
        present = strategyConditions[0]
        mergedStrategyConditions.append(present)
        strategyConditions.remove(present)
        temp = []
        for s in strategyConditions:
            if present.strategyId == s.strategyId \
                    and present.pickDate == s.pickDate \
                    and json.dumps(jsonable_encoder(present.strategyParams)) == json.dumps(
                jsonable_encoder(s.strategyParams)):
                temp.append(s)
                present.expressions.extend(s.expressions)
        # 找到一致的,添加完表达式就删除
        for t in temp:
            strategyConditions.remove(t)
    strategyIds = set([s.strategyId for s in mergedStrategyConditions])
    compilers = {}
    for strategyId in strategyIds:
        # 查询策略
        strategy = strategyDao.selectById(db, strategyId)
        assert_that(strategy).is_none().then_throw(ExceptionCode.STRATEGY_DOES_NOT_EXIST)
        strategyCodeQuery = StrategyCodeQuery(state=ConstantItem.STATE_ENABLED, strategyId=strategyId)
        strategyCode = strategyCodeDao.selectOne(db, strategyCodeQuery)
        assert_that(strategyCode).is_none().then_throw(ExceptionCode.STRATEGY_DOES_NOT_EXIST)
        context = CompilerContext(strategyCode.code)
        # context.strategyVarsCode = variables
        # context.pickResultCode = expressions
        compilers[strategyId] = TdxCompiler(context)
    # 设置K线查询参数
    adjust = ConstantItem.ADJUST_FORE if publicConditions.isAdjustFore else ConstantItem.ADJUST_NO
    isSt = 0 if publicConditions.isRemoveSt else None
    tradeStatus = 1 if publicConditions.isRemoveUntrade else None
    # pickDate最大日期不是st的股票
    pageParam = PageParam(pageNum=1, pageSize=500)
    for instrumentId in taskMsgBody.instrumentIds:
        try:
            tempResult = []
            for cond in mergedStrategyConditions:
                # 装载数据
                candlestickQuery = CandlestickQuery(instrumentId=instrumentId,
                                                    frequency=picker.frequency,
                                                    adjust=adjust,
                                                    tradeStatus=tradeStatus,
                                                    barDate=DateUtils.addDays(DateUtils.today(), 1))
                dbData = barDao.getCandlestick(db, candlestickQuery, pageParam)
                instrument = instrumentDao.selectById(db, instrumentId)
                assert_that(instrument).is_none().then_throw(ExceptionCode.INSTRUMENT_DOES_NOT_EXIST)
                candlestickQuery.symbol = instrument.symbol
                lastBars = redisManager.getLastCandlestick(candlestickQuery)
                dbData.reverse()
                if len(dbData) == 0:
                    dbData = dbData + lastBars
                else:
                    for bar in lastBars:
                        if bar.barDate > dbData[-1].barDate:
                            dbData.append(bar)
                # 装载数据
                candlestick = Candlestick()
                for row in dbData:
                    candlestick.closeSeries.append(row.close)
                    candlestick.openSeries.append(row.open)
                    candlestick.highSeries.append(row.high)
                    candlestick.lowSeries.append(row.low)
                    candlestick.dateSeries.append(DateUtils.dateToTimestamp13(row.barDate))
                    candlestick.volumeSeries.append(row.volume)
                    candlestick.amountSeries.append(row.amount)
                    candlestick.preCloseSeries.append(row.preClose)
                CLOSE.series = np.array(candlestick.closeSeries).astype(np.float)
                HIGH.series = np.array(candlestick.highSeries).astype(np.float)
                OPEN.series = np.array([candlestick.openSeries]).astype(np.float)
                LOW.series = np.array([candlestick.lowSeries]).astype(np.float)
                VOL.series = np.array([candlestick.volumeSeries]).astype(np.float)
                AMOUNT.series = np.array([candlestick.amountSeries]).astype(np.float)
                DATETIME.series = np.array(candlestick.dateSeries)
                # 初始化执行器
                # 转换策略参数
                variables = ""
                for param in cond.strategyParams:
                    variables += "%s = %s;\n" % (param.variable, param.value)
                # 转换选股结果表达式
                expressions = ""
                for expr in cond.expressions:
                    expressions += f" & {expr.parsedExpr}" if expressions else expr.parsedExpr

                # 编译源码
                compiler = compilers.get(cond.strategyId)
                context = compiler.context
                context.strategyVarsCode = variables
                context.pickResultCode = f'{context.varResult} = {expressions};'
                compiler.analyze().parse().generate().execute()
                # 加载数据
                tempResult.append(context.executeResult[context.varResult].series[-1])
                # drawings = executor.get_drawings()
            result = None
            for r in tempResult:
                if not result:
                    result = r
                else:
                    result &= r
            # 保存
            pickerDao.updatePickResult(db, picker.id, result)
            if result:
                pickerInstrument = PickerInstrument(pickerId=picker.id,
                                                    instrumentId=instrumentId,
                                                    probability=1)
                db.add(pickerInstrument)
            db.commit()
        except Exception as ex:
            print(f"{instrumentId}执行失败")
            pickerDao.updatePickResult(db, picker.id, False)
            db.commit()
            print(ex)
    db.close()
    return f"success"
